import cv2 as cv
import sys
import os
from python_ai.common.xcommon import *
import matplotlib.pyplot as plt
import numpy as np
import time
import datetime


def my_show_img(img, title="no title", trans=None, **kwargs):
    global spn
    spn += 1
    plt.subplot(spr, spc, spn)
    if trans is not None:
        img = trans(img)
    plt.imshow(img, **kwargs)
    plt.axis('off')
    plt.title(title)


# img_dir = '../../../../../large_data/CV2/lesson/Day06'
# img_name = 'hammer.jpg'
img_dir = '../../../../../large_data/pic'
# img_name = 'DSC05039.JPG'
# img_name = 'DSC05022_1.JPG'
# img_name = 'dog.jpg'
# img_name = 'dog_bird.jpg'
# img_name = 'football.jpg'
img_name = 'messi5.jpg'
img_path = os.path.join(img_dir, img_name)

spr = 2
spc = 3
spn = 0
plt.figure(figsize=[9, 6])

sep('load')
img = cv.imread(img_path, cv.IMREAD_GRAYSCALE)
# img = cv.resize(img, None, fx=0.1, fy=0.1, interpolation=cv.INTER_CUBIC)
print('original shape', img.shape)
H, W = img.shape
H2 = (H - 1) // 2
W2 = (W - 1) // 2
my_show_img(img, 'ori opencv', cmap='gray')

img = img.astype(np.float32)  # ATTENTION cv.dft needs float
f = cv.dft(img, flags=cv.DFT_COMPLEX_OUTPUT)  # ATTENTION complex_output
eps = 1e-8
f_show = np.log(cv.magnitude(f[:, :, 0], f[:, :, 1]) + eps)
my_show_img(f_show, 'f_show', cmap='gray')

f_shift = np.fft.fftshift(f)
f_shift_show = np.log(cv.magnitude(f_shift[:, :, 0], f_shift[:, :, 1]) + eps)
my_show_img(f_shift_show, 'f_shift_show', cmap='gray')

HALF = 30
mask = np.zeros((H, W, 2), dtype=np.float32)
mask[H2 - HALF:H2 + HALF, W2 - HALF:W2 + HALF] = 1.
f_shift *= mask
f_shift_show = np.log(cv.magnitude(f_shift[:, :, 0], f_shift[:, :, 1]) + eps)
my_show_img(f_shift_show, 'f_shift_show clipped', cmap='gray')

if_ = np.fft.ifftshift(f_shift)
if_show = np.log(cv.magnitude(if_[:, :, 0], if_[:, :, 1]) + eps)
my_show_img(if_show, 'if_show', cmap='gray')

img_ = cv.idft(if_)
img_ = cv.magnitude(img_[:, :, 0], img_[:, :, 1])
print_numpy_ndarray_info(img_, 'img_')
my_show_img(img_, 'filtered', cmap='gray')

plt.show()
